// Copyright 2009 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Author: Jeff Stewart
// Implements a read-only Registry access interface.

#include <windows.h>
#include "config/registrykey.h"

namespace popdiag {

namespace {

char hex_char[] = { '0', '1', '2', '3', '4', '5', '6', '7',
                    '8', '9', 'A', 'B', 'C', 'D', 'E', 'F' };
QString hex(unsigned char c) {
  QString rval;
  int lo = c & 0x0f;
  int hi = c >> 4;
  rval.append(hex_char[lo]);
  rval.append(hex_char[hi]);
  return rval;
}

// Yet another string conversion format.  This time, to and from std::wstring.
std::wstring QStringToWString(const QString &in) {
  std::wstring rval;
  for (int i = 0; i < in.length(); i++) {
    rval.push_back(static_cast<wchar_t>(in[i].unicode()));
  }
  return rval;
}

QString WStringToQString(const std::wstring &in) {
  QString rval;
  for (size_t i = 0; i < in.length(); i++) {
    rval.push_back(QChar(in[i]));
  }
  return rval;
}

}  // anonymous namespace

// Expand environment variables in the value represented
// as %xyz% using the windows ExpandEnvironmentStrings() call.
// Of course, lots of string conversion is needed, since
// the windows API operates on wstrings and we operate
// on QStrings.
QString ExpandEnvironmentReferences(const QString &in) {
  // convert to a wstring for manipulation:
  std::wstring qval = QStringToWString(in);

  // find out how much room we need
  DWORD expanded_size_needed
    = ExpandEnvironmentStrings(qval.data(), NULL, 0);

  WCHAR *buffer = new WCHAR[expanded_size_needed];
  ExpandEnvironmentStrings(qval.data(), buffer,
                           expanded_size_needed);

  // and back to QString-space, since we want to hide
  // wstring and its ilk from our callers:
  QString rval = WStringToQString(std::wstring(buffer));
  delete [] buffer;
  return rval;
}

///////////////////////////////////////////////////////////////////////
//  RegistryKeyPrivate Definition
///////////////////////////////////////////////////////////////////////

// Implementation of RegistryKey backed by the registry itself
class RegistryKeyPrivate : public RegistryKey {
 public:
  // Create an instance wrapping handle.  The destructor will
  // close the reference unless close_on_destroy is true, which
  // should only be set for the system root entries (CurrentUser,
  // etc.)
  RegistryKeyPrivate(HKEY handle, bool close_on_destroy);
  ~RegistryKeyPrivate();

  // All these methods match the corresponding RegistryKey versions.
  // See registrykey.h for descriptions:
  bool GetValueSizeInBytes(const QString &value_name, int *out) const;

  bool GetChildKeyNames(QList<QString> *out) const;
  RegistryKey *ChildKey(const QString &sub_key_name) const;

  bool ReadIntEntry(const QString &value_name, int *out) const;
  bool ReadStringEntry(const QString &value_name, QString *out) const;
  bool ReadBinaryString(const QString &value_name, QString *out) const;
  unsigned char *ReadBinaryEntry(const QString &value_name,
                                 size_t *out_len_in_bytes) const;

  QString Dump() const;
  RegistryKey *Copy() const;
  void Refresh();

 private:
  HKEY h_;
  bool close_;

  DISALLOW_EVIL_CONSTRUCTORS(RegistryKeyPrivate);
};

///////////////////////////////////////////////////////////////////////
//  RegistryKey Implementation
///////////////////////////////////////////////////////////////////////

RegistryKey::RegistryKey() {
}

RegistryKey::~RegistryKey() {
}

const RegistryKey &StaticLazyInitRegistryKey(
    HKEY handle, scoped_ptr<RegistryKey> *fill) {
  if (*fill == NULL) {
    fill->reset(new RegistryKeyPrivate(handle, false));
  }
  return *(fill->get());
}

static scoped_ptr<RegistryKey> kCurrentUser;
static scoped_ptr<RegistryKey> kClassesRoot;
static scoped_ptr<RegistryKey> kLocalMachine;

const RegistryKey &RegistryKey::CurrentUser() {
  return StaticLazyInitRegistryKey(HKEY_CURRENT_USER, &kCurrentUser);
}

const RegistryKey &RegistryKey::ClassesRoot() {
  return StaticLazyInitRegistryKey(HKEY_CLASSES_ROOT, &kClassesRoot);
}

const RegistryKey &RegistryKey::LocalMachine() {
  return StaticLazyInitRegistryKey(HKEY_LOCAL_MACHINE, &kLocalMachine);
}

///////////////////////////////////////////////////////////////////////
//  RegistryKeyPrivate Implementation
///////////////////////////////////////////////////////////////////////

RegistryKeyPrivate::RegistryKeyPrivate(HKEY handle, bool close)
  : RegistryKey(), h_(handle), close_(close) {
}

RegistryKeyPrivate::~RegistryKeyPrivate() {
  if (h_ != NULL && close_) {
    RegCloseKey(h_);
    h_ = NULL;
  }
}

RegistryKey *RegistryKeyPrivate::Copy() const {
  if (!close_) {
    return const_cast<RegistryKeyPrivate *>(this);
  }
  HKEY child = NULL;
  CHECK(ERROR_SUCCESS == RegOpenKeyEx(h_, NULL, 0, KEY_READ, &child));
  return new RegistryKeyPrivate(child, true);
}

RegistryKey *RegistryKeyPrivate::ChildKey(const QString &sub_key_name) const {
  HKEY child = NULL;
  std::wstring key_name_w(QStringToWString(sub_key_name));
  DWORD result = RegOpenKeyEx(h_, key_name_w.c_str(), 0, KEY_READ, &child);
  if (ERROR_SUCCESS == result) {
    return new RegistryKeyPrivate(child, true);
  }
  return NULL;
}

unsigned char *RegistryKeyPrivate::ReadBinaryEntry(const QString &value_name,
  size_t *out_len_in_bytes) const {
  DWORD type_found;
  int size_needed;
  std::wstring value_name_w(QStringToWString(value_name));

  if (!this->GetValueSizeInBytes(value_name, &size_needed)) {
    return NULL;
  }
  unsigned char *rval = new unsigned char[size_needed];
  DWORD size_found = size_needed;
  CHECK(ERROR_SUCCESS == RegQueryValueEx(h_, value_name_w.c_str(), NULL,
                                         &type_found, rval, &size_found));
  if (type_found != REG_BINARY) {
    delete [] rval;
    return NULL;
  }
  *out_len_in_bytes = size_found;
  return rval;
}

bool RegistryKeyPrivate::ReadBinaryString(const QString &value_name,
                                          QString *out) const {
  size_t len = 0;
  unsigned char *binary_form = ReadBinaryEntry(value_name, &len);
  if (binary_form == NULL) {
    return ReadStringEntry(value_name, out);
  }
  size_t len_in_wchar = len / sizeof(WCHAR);
  CHECK_GT(len_in_wchar, 0);
  WCHAR *in = reinterpret_cast<WCHAR *>(binary_form);
  CHECK_EQ(in[len_in_wchar - 1], 0);
  *out = WStringToQString(std::wstring(in));
  return true;
}


bool RegistryKeyPrivate::ReadStringEntry(const QString &value_name,
                                         QString *out) const {
  int size_needed;
  if (!this->GetValueSizeInBytes(value_name, &size_needed)) {
    return false;
  }

  std::wstring value_name_w(QStringToWString(value_name));
  scoped_ptr<WCHAR> buffer(new WCHAR[size_needed]);
  DWORD entry_type;
  DWORD size_found = size_needed;
  if (ERROR_SUCCESS ==
      RegQueryValueEx(h_, value_name_w.c_str(), NULL, &entry_type,
        reinterpret_cast<LPBYTE>(buffer.get()), &size_found)) {
    if (entry_type != REG_SZ && entry_type != REG_EXPAND_SZ) {
      return false;
    }
    *out = WStringToQString(std::wstring(buffer.get()));
    if (entry_type == REG_EXPAND_SZ) {
      *out = ExpandEnvironmentReferences(*out);
    }
    return true;
  }
  return false;
}

bool RegistryKeyPrivate::ReadIntEntry(const QString &value_name,
                                      int *out) const {
  DWORD size;
  DWORD entry_type;
  DWORD out_dword;
  std::wstring value_name_w(QStringToWString(value_name));
  HRESULT rval = RegQueryValueEx(h_, value_name_w.c_str(), NULL,
                                 &entry_type,
                                 reinterpret_cast<LPBYTE>(&out_dword),
                                 &size);
  if (!SUCCEEDED(rval)) {
    return false;
  }
  if (entry_type != REG_DWORD) {
    return false;
  }
  *out = out_dword;
  return true;
}

bool RegistryKeyPrivate::GetChildKeyNames(QList<QString> *out) const {
  out->clear();
  DWORD num_sub_keys = 0;
  if (ERROR_SUCCESS != RegQueryInfoKey(h_, NULL, NULL, NULL, &num_sub_keys,
    NULL, NULL, NULL, NULL, NULL, NULL, NULL)) {
    return false;
  }
  // key names max out at 255.
  WCHAR buffer[256];
  for (DWORD i = 0; i < num_sub_keys; i++) {
    DWORD buffer_size = sizeof(buffer) / sizeof(TCHAR);
    HRESULT rval = RegEnumKeyEx(h_, i, buffer, &buffer_size,
                                NULL, NULL, NULL, NULL);
    if (rval == ERROR_NO_MORE_ITEMS) {
      return true;
    }
    if (rval != ERROR_SUCCESS) {
      return false;
    }
    QString val(WStringToQString(std::wstring(buffer)));
    out->push_back(val);
  }
  return true;
}

QString RegistryKeyPrivate::Dump() const {
  int i = 0;
  QString str;
  QTextStream log(&str, QIODevice::WriteOnly);
  log << "<table>";
  while (true) {
    WCHAR elem_name[256];
    DWORD name_size = sizeof(elem_name) / sizeof(TCHAR);
    DWORD elem_type;
    BYTE data_value[512];
    DWORD data_size = sizeof(data_value);
    HRESULT hr = RegEnumValue(h_, i++, elem_name, &name_size, NULL, &elem_type,
                              data_value, &data_size);
    if (hr == ERROR_NO_MORE_ITEMS) {
      break;
    }
    if (FAILED(hr)) {
      break;
    }
    log << "<tr><td>";
    log << QString::fromUtf16(reinterpret_cast<ushort *>(elem_name))
        << "</td><td>";
    switch (elem_type) {
      case REG_DWORD:
        log << "REG_DWORD</td><td>" << *reinterpret_cast<DWORD *>(data_value)
            << "</td></tr>";
        break;
      case REG_SZ:
        // verify null termination:
        CHECK(data_size <= sizeof(data_value));
        data_value[data_size] = '\0';
        log << "REG_SZ</td><td>\""
            << QString::fromUtf16(reinterpret_cast<ushort *>(data_value))
            << "\"</td></tr>";
        break;
      case REG_BINARY:
        {
          log << "REG_BINARY</td><td>";
          for (DWORD i = 0; i < data_size; i++) {
            if ((i % 8) == 0) {
              log << " ";
            }
            log << hex(data_value[i]);
          }
          // because so many of the items we're interested in
          // are strings stored as binary, try to print it as
          // a string:
          if (data_size % 2 == 0) {
            log << " :: ";
            log << dec;
            ushort *data = reinterpret_cast<ushort *>(data_value);
            int ushort_size = data_size / 2;
            for (int i = 0; i < ushort_size; i++) {
              QChar c(data[i]);
              if (c.isPrint()) {
                log << c;
              } else {
                log << ".";
              }
            }
          }
          log << "</td></tr>";
        }
        break;
      case REG_DWORD_BIG_ENDIAN:
      case REG_EXPAND_SZ:
      case REG_LINK:
      case REG_MULTI_SZ:
      case REG_NONE:
      case REG_QWORD:
      default:
        log << elem_type << "</td><td>not printed</td></tr>";
    }
    log << "\n";
  }
  log << "</table>\n";
  return str;
}

bool RegistryKeyPrivate::GetValueSizeInBytes(const QString &value_name,
                                             int *out) const {
  DWORD type_found;
  DWORD size_found;
  std::wstring value_name_w(QStringToWString(value_name));
  if (ERROR_SUCCESS != RegQueryValueEx(h_, value_name_w.c_str(), NULL,
                                       &type_found, NULL, &size_found)) {
    return false;
  }
  *out = size_found;
  return true;
}

// Always reads current data.  No Refresh needed.
void RegistryKeyPrivate::Refresh() {
}

}  // namespace popdiag
